# simple_arm_env.py
from __future__ import annotations
import os
import time
from typing import Optional, Tuple

import gymnasium as gym
from gymnasium import spaces
import numpy as np
import pybullet as p
import pybullet_data


class SimpleArmEnv(gym.Env):
    """
    一个将你的 PyBullet 机械臂封装为 Gymnasium 环境的最小可用版本。
    关节0：力矩控制（动作 ∈ [-1,1] 缩放为 [-max_torque, +max_torque]）
    关节1：被动（仅受 URDF 阻尼/摩擦/重力）
    观测： [q0, q1, dq0, dq1, ee_x, ee_y, ee_z] 共7维
    """
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 240}

    def __init__(
        self,
        urdf_path: str = "./urdf/simple_arm.urdf",
        render_mode: Optional[str] = None,
        max_torque: float = 50.0,
        time_step: float = 1.0 / 100.0,
        frame_skip: int = 1,
        max_episode_steps: int = 2500,
        gui_camera_target=(0.0, 0.0, 1.05),
        gui_camera_distance: float = 2.0,
        gui_camera_yaw: float = 45.0,
        gui_camera_pitch: float = -30.0,
    ):
        super().__init__()
        assert render_mode in (None, "human", "rgb_array")
        self.max_episode_steps = max_episode_steps
        self.current_step = 0  # 当前步数计数器
        self.render_mode = render_mode
        self.urdf_path = urdf_path
        self.max_torque = float(max_torque)
        self.time_step = float(time_step)
        self.frame_skip = int(frame_skip)

        # ---- Spaces ----
        self.action_space = spaces.Box(
            low=np.array([-1.0], dtype=np.float32),
            high=np.array([+1.0], dtype=np.float32),
            dtype=np.float32,
        )
        high = np.array([np.inf, np.inf, np.inf, np.inf], dtype=np.float32)
        self.observation_space = spaces.Box(-high, +high, dtype=np.float32)

        # ---- Bullet connection ----
        self._cid: int = p.connect(p.GUI if self.render_mode == "human" else p.DIRECT)
        p.setTimeStep(self.time_step)
        p.setGravity(0, 0, -9.8)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())

        # 以 -1 作为“未创建”的无效 id，避免 int|None 的类型告警
        self.plane_id: int = -1
        self.robot_id: int = -1

        # 关节索引（根据你的 URDF：0=base_joint, 1=elbow_joint）
        self.joint_indices = (0, 1)

        # GUI 相机参数（仅 GUI 下生效）
        self.gui_camera_target = tuple(gui_camera_target)
        self.gui_camera_distance = float(gui_camera_distance)
        self.gui_camera_yaw = float(gui_camera_yaw)
        self.gui_camera_pitch = float(gui_camera_pitch)

        self._build_world()

    # ---------- World building ----------
    def _build_world(self) -> None:
        if self.plane_id != -1:
            p.removeBody(self.plane_id)
            self.plane_id = -1
        if self.robot_id != -1:
            p.removeBody(self.robot_id)
            self.robot_id = -1

        plane_shape = p.createCollisionShape(p.GEOM_PLANE)
        self.plane_id = p.createMultiBody(baseMass=0, baseCollisionShapeIndex=plane_shape)

        if not os.path.exists(self.urdf_path):
            raise FileNotFoundError(f"URDF not found: {self.urdf_path}")

        self.robot_id = p.loadURDF(self.urdf_path, [0, 0, 0.001], useFixedBase=False)

        # 释放关节（关闭默认速度电机），之后才好用 TORQUE_CONTROL
        for j in self.joint_indices:
            p.setJointMotorControl2(self.robot_id, j, p.VELOCITY_CONTROL, force=0.0)

        # 轻微固定基座，避免数值漂移
        p.createConstraint(
            parentBodyUniqueId=self.robot_id,
            parentLinkIndex=-1,
            childBodyUniqueId=-1,
            childLinkIndex=-1,
            jointType=p.JOINT_FIXED,
            jointAxis=[0, 0, 0],
            parentFramePosition=[0, 0, -0.0005],
            childFramePosition=list(self.gui_camera_target),
        )

        if self.render_mode == "human":
            p.resetDebugVisualizerCamera(
                cameraDistance=self.gui_camera_distance,
                cameraYaw=self.gui_camera_yaw,
                cameraPitch=self.gui_camera_pitch,
                cameraTargetPosition=list(self.gui_camera_target),
            )

    # ---------- Helpers ----------
    def _get_joint_state(self) -> Tuple[float, float, float, float]:
        assert self.robot_id != -1, "robot not loaded"
        s0 = p.getJointState(self.robot_id, self.joint_indices[0])
        s1 = p.getJointState(self.robot_id, self.joint_indices[1])
        q0, dq0 = float(s0[0]), float(s0[1])
        q1, dq1 = float(s1[0]), float(s1[1])
        return q0, q1, dq0, dq1

    def _get_ee_pos(self) -> Tuple[float, float, float]:
        """末端位置：优先取 child link 名为 'ball' 的那个 link；否则退化为 elbow 的 child link。"""
        assert self.robot_id != -1, "robot not loaded"
        ee_link_index = None
        n = p.getNumJoints(self.robot_id)
        for i in range(n):
            # jointInfo[12] 是 child link 名称
            child_link_name = p.getJointInfo(self.robot_id, i)[12].decode("utf-8")
            if child_link_name == "ball":
                ee_link_index = i
                break
        if ee_link_index is None:
            ee_link_index = self.joint_indices[-1]

        pos = p.getLinkState(self.robot_id, ee_link_index, computeForwardKinematics=True)[4]
        return float(pos[0]), float(pos[1]), float(pos[2])

    def _get_obs(self) -> np.ndarray:
        q0, q1, dq0, dq1 = self._get_joint_state()
        # ex, ey, ez = self._get_ee_pos()
        return np.array([q0, q1, dq0, dq1], dtype=np.float32)

    # ---------- Gymnasium API ----------
    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
        # Gymnasium 规范：用 super().reset(seed=seed) 初始化内部 RNG
        super().reset(seed=seed)
        
        self.current_step = 0  # 重置步数计数器
        p.resetSimulation()
        p.setTimeStep(self.time_step)
        p.setGravity(0, 0, -10)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())

        self.plane_id = -1
        self.robot_id = -1
        self._build_world()

        # 轻微随机初始角度
        rng = np.random.default_rng(seed)
        for j in self.joint_indices:
            angle = float(rng.uniform(-0.05, 0.05))
            p.resetJointState(self.robot_id, j, targetValue=angle, targetVelocity=0.0)

        obs = self._get_obs()
        return obs, {}

    def step(self, action: np.ndarray):
        assert self.robot_id != -1, "robot not loaded"

        self.current_step += 1

        a = float(np.clip(action[0], -1.0, 1.0))
        tau = a * self.max_torque

        # 关节0：力矩控制（关节1保持被动）
        p.setJointMotorControl2(self.robot_id, self.joint_indices[0], p.TORQUE_CONTROL, force=tau)

        for _ in range(self.frame_skip):
            p.stepSimulation()

        obs = self._get_obs()

        # 示例奖励（可按需替换）
        z = float(obs[-1])
        q0, q1, dq0, dq1 = map(float, obs[:4])
        reward = -abs(q1)
        # print(f"q0={q0:.3f}, q1={q1:.3f}, dq0={dq0:.3f}, dq1={dq1:.3f}")
        terminated = False
        truncated = False
        if not np.isfinite(obs).all() or abs(q0) > 100 or abs(q1) > 100 or self.current_step >= self.max_episode_steps:
            truncated = True

        info = {}
        if self.render_mode == "human":
            self.render()
        return obs, float(reward), terminated, truncated, info

    def render(self):
        if self.render_mode == "rgb_array":
            width, height = 640, 480
            view_matrix = p.computeViewMatrixFromYawPitchRoll(
                cameraTargetPosition=list(self.gui_camera_target),
                distance=self.gui_camera_distance,
                yaw=self.gui_camera_yaw,
                pitch=self.gui_camera_pitch,
                roll=0,
                upAxisIndex=2,
            )
            proj_matrix = p.computeProjectionMatrixFOV(
                fov=60, aspect=float(width) / float(height), nearVal=0.01, farVal=5.0
            )
            img = p.getCameraImage(width, height, view_matrix, proj_matrix)[2]
            return np.asarray(img, dtype=np.uint8)
        elif self.render_mode == "human":
            return None

    def close(self):
        if p.isConnected():
            p.disconnect()

if __name__ == "__main__":
    env = SimpleArmEnv(render_mode="human")
    obs, info = env.reset()
    for _ in range(100000):
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        # print(f"obs={obs}, reward={reward}, terminated={terminated}, truncated={truncated}")
        time.sleep(1.0 / 100.0)
        if terminated or truncated:
            obs, info = env.reset()
    env.close()